import cv2
import os
import numpy as np
import tkinter as tk
from tkinter import messagebox, simpledialog, filedialog
import datetime
import threading
import face_recognition  # 安装 face_recognition 库

# 创建一个保存人脸图像的目录
database_dir = "face_database"
if not os.path.exists(database_dir):
    os.makedirs(database_dir)

# 创建一个导入数据集的目录
import_dataset_dir = "import_dataset"
if not os.path.exists(import_dataset_dir):
    os.makedirs(import_dataset_dir)

# 初始化 HAAR 分类器
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

# 存储人脸图像及其名称
known_faces = {}
face_encodings = []
capture_active = False
current_face_image = None



def load_known_faces(img_path):
    # 打印加载的图像路径
    print("Loading image from:", img_path)

    # 加载图像
    image = face_recognition.load_image_file(img_path)

    # 获取图像中的人脸编码
    encodings = face_recognition.face_encodings(image)

    # 检查是否找到了任何人脸编码
    if len(encodings) == 0:
        print(f"No faces found in the image: {img_path}")
        return None  # 或者其他处理方式，如返回空列表

    # 如果找到了人脸编码，获取第一个编码
    encoding = encodings[0]
    return encoding

# 使用绝对路径来加载图像
img_path = "C:\\Users\\long\\PycharmProjects\\face_recognition_long\\face_database\\Male_140075517484900.jpg"  # 替换为实际图像路径
known_face_encoding = load_known_faces(img_path)

if known_face_encoding is not None:
    print("Face encoding loaded successfully.")
else:
    print("Failed to load face encoding.")

def classify_gender(face_image):
    # 这里应该是您模型的预测逻辑
    return np.random.choice(['Male', 'Female'])  # 仅作为示例，您需要使用实际模型


def start_recognition():
    global capture_active, current_face_image
    capture_active = True
    cap = cv2.VideoCapture(0)

    if not cap.isOpened():
        messagebox.showerror("错误", "无法打开摄像头！")
        return

    while capture_active:
        ret, frame = cap.read()
        if not ret:
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)

        # 初始化 name 变量
        name = None

        # 处理检测到的每个人脸
        if len(faces) > 0:
            for (x, y, w, h) in faces:
                face_image = frame[y:y + h, x:x + w]
                rgb_face_image = cv2.cvtColor(face_image, cv2.COLOR_BGR2RGB)  # 转为RGB格式
                face_encodings_in_frame = face_recognition.face_encodings(rgb_face_image)

                if face_encodings_in_frame:  # 检测到人脸编码
                    current_face_encoding = face_encodings_in_frame[0]

                    # 在数据库中寻找匹配
                    matches = face_recognition.compare_faces(face_encodings, current_face_encoding)
                    name = None

                    if True in matches:
                        first_match_index = matches.index(True)
                        name = list(known_faces.keys())[first_match_index]  # 获取匹配的名字

                    # 绘制矩形框和显示名字
                    if name:  # 如果找到匹配的居民
                        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)  # 绿色矩形框
                        cv2.putText(frame, name, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                    else:  # 如果是陌生人
                        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 255), 2)  # 红色矩形框

                # 仅在用户点击住户命名按钮时才保存人脸图像
                if capture_active and (name is None):  # 只有陌生人时才保存
                    current_face_image = face_image  # 保存当前人脸图像
                    break

        cv2.imshow("Face Recognition", frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


def resident_name():
    global current_face_image
    if current_face_image is not None:
        name = simpledialog.askstring("命名", "请输入住户姓名：", parent=root)
        if name:
            gender = classify_gender(current_face_image)
            filename = f"{name}_{gender}_{cv2.getTickCount()}.jpg"
            face_filename = os.path.join(database_dir, filename)
            cv2.imwrite(face_filename, current_face_image)
            known_faces[name] = filename  # 更新已知人脸字典
            face_encodings.append(face_recognition.face_encodings(face_recognition.load_image_file(face_filename))[0])  # 添加编码
            messagebox.showinfo("成功", f"已保存人脸图像：{face_filename}")
            current_face_image = None
        else:
            messagebox.showwarning("警告", "请提供有效的住户姓名！")
    else:
        messagebox.showwarning("警告", "现在没有可保存的人脸图像！")

def rename_image():
    file_path = filedialog.askopenfilename(initialdir=database_dir, title="选择要重命名的文件",
                                           filetypes=(("JPEG files", "*.jpg"), ("all files", "*.*")))
    if not file_path:
        return

    new_name = simpledialog.askstring("重命名", "请输入新的文件名（不带扩展名）：")
    if new_name:
        new_path = os.path.join(database_dir, f"{new_name}.jpg")
        os.rename(file_path, new_path)
        messagebox.showinfo("成功", f"文件已重命名为：{new_path}")
    else:
        messagebox.showwarning("警告", "请提供有效的文件名！")

def exit_application():
    if messagebox.askokcancel("退出确认", "您确定要退出吗？"):
        global capture_active
        capture_active = False
        root.quit()

def import_and_rename_images():
    base_directory = "import_dataset"  # 数据集的根目录

    # 遍历公司文件夹
    for company_folder in os.listdir(base_directory):
        company_path = os.path.join(base_directory, company_folder)

        # 检查文件夹是否存在且是一个目录
        if os.path.isdir(company_path):
            for idx, image_file in enumerate(os.listdir(company_path), start=1):
                if image_file.endswith('.bmp'):  # 只处理bmp格式的图片
                    image_path = os.path.join(company_path, image_file)
                    face_image = cv2.imread(image_path)

                    if face_image is not None:
                        # 进行性别识别
                        gender = classify_gender(face_image)

                        # 生成时间戳
                        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

                        # 生成新文件名
                        new_filename = f"{company_folder}_{gender}_{idx}_{timestamp}.jpg"
                        save_path = os.path.join(database_dir, new_filename)

                        # 保存重命名后的图像
                        cv2.imwrite(save_path, face_image)
                        messagebox.showinfo("成功", f"已保存：{save_path}")

def toggle_capture():
    global capture_active
    capture_active = not capture_active
    state = "开始" if capture_active else "停止"
    messagebox.showinfo("状态", f"人脸采集已{state}。")

def login():
    login_window = tk.Toplevel(root)
    login_window.title("登录")

    tk.Label(login_window, text="用户名:").grid(row=0, column=0)
    tk.Label(login_window, text="密码:").grid(row=1, column=0)

    username_entry = tk.Entry(login_window)
    password_entry = tk.Entry(login_window, show='*')

    username_entry.grid(row=0, column=1)
    password_entry.grid(row=1, column=1)

    def validate_login():
        if username_entry.get() == "admin" and password_entry.get() == "admin":
            login_window.destroy()
            main_interface()
        else:
            messagebox.showerror("错误", "用户名或密码错误！")

    tk.Button(login_window, text="登录", command=validate_login).grid(row=2, columnspan=2)

def main_interface():
    start_button = tk.Button(root, text="启动人脸识别", command=start_recognition_thread, font=('Arial', 16))
    start_button.pack(pady=20)

    resident_name_button = tk.Button(root, text="住户命名", command=resident_name, font=('Arial', 16))
    resident_name_button.pack(pady=10)

    rename_button = tk.Button(root, text="重命名人脸图像", command=rename_image, font=('Arial', 16))
    rename_button.pack(pady=10)

    import_button = tk.Button(root, text="导入并重命名人脸图像", command=import_and_rename_images, font=('Arial', 16))
    import_button.pack(pady=10)

    toggle_button = tk.Button(root, text="开始/停止 用户命名", command=toggle_capture, font=('Arial', 16))
    toggle_button.pack(pady=10)

    exit_button = tk.Button(root, text="退出程序", command=exit_application, font=('Arial', 16))
    exit_button.pack(pady=10)

def start_recognition_thread():
    recognition_thread = threading.Thread(target=start_recognition)
    recognition_thread.start()

# 创建主 GUI 界面
root = tk.Tk()
root.title("人脸识别系统")
root.geometry("400x500")

# 调用登录函数
login()

root.mainloop()
